{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flocks, Herds, and Traffic Jams\n",
    "\n",
    "Code examples from [Think Complexity, 2nd edition](http://greenteapress.com/wp/complexity2), Chapter 10\n",
    "\n",
    "Copyright 2016 Allen Downey, [MIT License](http://opensource.org/licenses/MIT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import vpython\n",
    "\n",
    "from __future__ import print_function, division\n",
    "\n",
    "%matplotlib inline\n",
    "%precision 3\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Boids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# size of the boids\n",
    "b_radius = 0.03\n",
    "b_length = 0.1\n",
    "\n",
    "# radiuses for sensing different rules\n",
    "r_avoid = 0.3\n",
    "r_center = 1.0\n",
    "r_copy = 0.5\n",
    "\n",
    "# viewing angle for different rules, in radians\n",
    "a_avoid = 2*np.pi\n",
    "a_center = 2\n",
    "a_copy = 2\n",
    "\n",
    "# weights for various rules\n",
    "w_avoid = 4\n",
    "w_center = 3\n",
    "w_copy = 2\n",
    "w_love = 10\n",
    "\n",
    "# time step\n",
    "dt = 0.1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def random_vector(a, b):\n",
    "    \"\"\"Create a vector with each element uniformly distributed in [a, b).\"\"\"\n",
    "    t = [np.random.uniform(a,b) for i in range(3)]\n",
    "    return vpython.vector(*t)\n",
    "\n",
    "\n",
    "def limit_vector(vect):\n",
    "    \"\"\"if the magnitude is greater than 1, set it to 1\"\"\"\n",
    "    if vect.mag > 1:\n",
    "        vect.mag = 1\n",
    "    return vect\n",
    "\n",
    "\n",
    "null_vector = vpython.vector(0,0,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<2.590807, 0.387244, 4.465229>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "random_vector(0, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Boid(vpython.cone):\n",
    "    \"\"\"A Boid is a VPython cone with a velocity\"\"\"\n",
    "\n",
    "    def __init__(self, radius=b_radius, length=b_length):\n",
    "        pos = random_vector(0, 1)\n",
    "        self.vel = random_vector(0, 1).norm()\n",
    "        vpython.cone.__init__(self, pos=pos, radius=radius)\n",
    "        self.axis = length * self.vel.norm()\n",
    "\n",
    "    def get_neighbors(self, others, radius, angle):\n",
    "        \"\"\"Return the list of neighbors within the given radius and angle.\"\"\"\n",
    "        boids = []\n",
    "        for other in others:\n",
    "            if other is self: continue\n",
    "            offset = other.pos - self.pos\n",
    "            \n",
    "            # if not in range, skip it\n",
    "            if offset.mag > radius: \n",
    "                continue\n",
    "\n",
    "            # if not within viewing angle, skip it\n",
    "            if self.vel.diff_angle(offset) > angle:\n",
    "                continue\n",
    "\n",
    "            # otherwise add it to the list\n",
    "            boids.append(other)\n",
    "            \n",
    "        return boids\n",
    "\n",
    "    def avoid(self, others, carrot):\n",
    "        \"\"\"Find the center of mass of all objects in range and\n",
    "        returns a vector in the opposite direction, with magnitude\n",
    "        proportional to the inverse of the distance (up to a limit).\"\"\"\n",
    "        others = others + [carrot]\n",
    "        close = self.get_neighbors(others, r_avoid, a_avoid)\n",
    "        t = [other.pos for other in close]\n",
    "        if t:\n",
    "            center = np.sum(t)/len(t)\n",
    "            away = vpython.vector(self.pos - center)\n",
    "            away.mag = r_avoid / away.mag\n",
    "            return limit_vector(away)\n",
    "        else:\n",
    "            return null_vector\n",
    "\n",
    "    def center(self, others):\n",
    "        \"\"\"Find the center of mass of other boids in range and\n",
    "        returns a vector pointing toward it.\"\"\"\n",
    "        close = self.get_neighbors(others, r_center, a_center)\n",
    "        t = [other.pos for other in close]\n",
    "        if t:\n",
    "            center = np.sum(t)/len(t)\n",
    "            toward = vpython.vector(center - self.pos)\n",
    "            return limit_vector(toward)\n",
    "        else:\n",
    "            return null_vector\n",
    "\n",
    "    def copy(self, others):\n",
    "        \"\"\"Return the average heading of other boids in range.\n",
    "        \n",
    "        others: list of Boids\n",
    "        \"\"\"\n",
    "        close = self.get_neighbors(others, r_copy, a_copy)\n",
    "        t = [other.vel for other in close]\n",
    "        if t:\n",
    "            # TODO: replace this with mean\n",
    "            center = np.sum(t)/len(t)\n",
    "            away = vpython.vector(self.pos - center)\n",
    "            return limit_vector(away)\n",
    "        else:\n",
    "            return null_vector\n",
    "\n",
    "    def love(self, carrot):\n",
    "        \"\"\"Returns a vector pointing toward the carrot.\"\"\"\n",
    "        toward = carrot.pos - self.pos\n",
    "        return limit_vector(toward)\n",
    "\n",
    "    def set_goal(self, boids, carrot):\n",
    "        \"\"\"Sets the goal to be the weighted sum of the goal vectors.\"\"\"\n",
    "        self.goal = (w_avoid * self.avoid(boids, carrot) + \n",
    "                     w_center * self.center(boids) +\n",
    "                     w_copy * self.copy(boids) + \n",
    "                     w_love * self.love(carrot))\n",
    "        self.goal.mag = 1\n",
    "        \n",
    "    def move(self, mu=0.1):\n",
    "        \"\"\"Update the velocity, position and axis vectors.\n",
    "        mu controls how fast the boids can turn (maneuverability).\"\"\"\n",
    "\n",
    "        self.vel = (1-mu) * self.vel + mu * self.goal\n",
    "        self.vel.mag = 1\n",
    "\n",
    "        self.pos += dt * self.vel\n",
    "        self.axis = b_length * self.vel.norm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class World(object):\n",
    "    \n",
    "    def __init__(self, n=10):\n",
    "        \"\"\"Create n Boids and one carrot.\n",
    "\n",
    "        tracking: indicates whether the carrot follows the mouse\n",
    "        \"\"\"\n",
    "        self.boids = [Boid() for i in range(n)]\n",
    "        vec = vpython.vector\n",
    "        self.carrot = vpython.sphere(pos=vec(1,0,0), radius=0.1, color=vec(1,0,0))\n",
    "        print(self.carrot)\n",
    "        self.tracking = False\n",
    "        \n",
    "    def step(self):\n",
    "        \"\"\"Compute one time step.\"\"\"\n",
    "        # move the boids\n",
    "        for boid in self.boids:\n",
    "            boid.set_goal(self.boids, self.carrot)\n",
    "            boid.move()\n",
    "\n",
    "        # mouse click toggles tracking\n",
    "        #if canvas.mouse.clicked:\n",
    "        #    canvas.mouse.getclick()\n",
    "        #    self.tracking = not self.tracking\n",
    "\n",
    "        # if we're tracking, move the carrot\n",
    "        if self.tracking:\n",
    "            self.carrot.pos = canvas.mouse.pos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div id=\"glowscript\" class=\"glowscript\"></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/javascript": [
       "window.__context = { glowscript_container: $(\"#glowscript\").removeAttr(\"id\")}"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<vpython.vpython.canvas object at 0x7f66440f0748>\n",
      "<vpython.vpython.sphere object at 0x7f6644132860>\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-9-f20a792385dd>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     12\u001b[0m \u001b[1;32mwhile\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     13\u001b[0m     \u001b[1;31m# update the screen once per time step\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 14\u001b[1;33m     \u001b[0mvpython\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrate\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m/\u001b[0m\u001b[0mdt\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     15\u001b[0m     \u001b[0mworld\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m/home/downey/anaconda2/envs/py3k/lib/python3.5/site-packages/vpython/vpython.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, maxRate)\u001b[0m\n\u001b[0;32m    109\u001b[0m         \u001b[1;32mif\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrval\u001b[0m \u001b[1;33m!=\u001b[0m \u001b[0mmaxRate\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mand\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mmaxRate\u001b[0m \u001b[1;33m>=\u001b[0m \u001b[1;36m1.0\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    110\u001b[0m             \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrval\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmaxRate\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 111\u001b[1;33m         \u001b[0msuper\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mRateKeeper2\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m__call__\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmaxRate\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    112\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    113\u001b[0m \u001b[1;32mif\u001b[0m \u001b[0msys\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mversion\u001b[0m \u001b[1;33m>\u001b[0m \u001b[1;34m'3'\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m/home/downey/anaconda2/envs/py3k/lib/python3.5/site-packages/vpython/rate_control.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, maxRate)\u001b[0m\n\u001b[0;32m    198\u001b[0m                 \u001b[1;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mwhenToRender\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrenderIndex\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m==\u001b[0m \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrateCount\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    199\u001b[0m                     \u001b[0msleeps\u001b[0m \u001b[1;33m+=\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 200\u001b[1;33m                     \u001b[0m_sleep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0minteractionPeriod\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    201\u001b[0m                 \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    202\u001b[0m                     \u001b[1;32mbreak\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m/home/downey/anaconda2/envs/py3k/lib/python3.5/site-packages/vpython/rate_control.py\u001b[0m in \u001b[0;36m_sleep\u001b[1;34m(dt)\u001b[0m\n\u001b[0;32m     47\u001b[0m         \u001b[0mdtsleep\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnticks\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0m_tick\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     48\u001b[0m         \u001b[0mt\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 49\u001b[1;33m         \u001b[0mtime\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msleep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdtsleep\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     50\u001b[0m         \u001b[0mt\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_clock\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m-\u001b[0m\u001b[0mt\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     51\u001b[0m         \u001b[0mdt\u001b[0m \u001b[1;33m-=\u001b[0m \u001b[0mt\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "n = 20\n",
    "size = 5\n",
    "\n",
    "canvas = vpython.canvas(title='Boids', width=800, height=600,\n",
    "                        range=(size, size, size))\n",
    "print(canvas)\n",
    "\n",
    "world = World(n)\n",
    "canvas.center = world.carrot.pos\n",
    "canvas.autoscale = False\n",
    "\n",
    "while 1:\n",
    "    # update the screen once per time step\n",
    "    vpython.rate(1/dt)\n",
    "    world.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
